keras 如何优雅地实现多个自定义损失函数 |
您所在的位置：网站首页 › keras 自定义优化 › keras 如何优雅地实现多个自定义损失函数 |
keras 如何优雅地实现多个自定义损失函数
|
keras 如何优雅地实现多个自定义损失函数
之前跑SRGAN模型的时候，由于有多个损失函数（content loss、adversarial loss、perceptual loss）当时采用的是train_on_batch进行训练的。但是发现相比于model.fit，train_on_batch的训练速度会慢很多（虽然我也不知道为啥）而且还无法显示loss变化的进度调，所以就让人很难受QAQ。后来看见知乎上这位作者的文章Tensorflow2.0中复杂损失函数实现，那现在就说一下我是怎么实现多个自定义损失函数的吧。 主要思想是把自定义的损失函数作为一个layer添加到模型中，然后再把自定义的损失作为output输出 以下代码基于tensorflow-gpu==2.5.0完成 首先还是先导包： import tensorflow as tf from tensorflow.keras import backend as K from tensorflow.keras import layers from tensorflow.keras import Model使用Subclass定义损失函数层（这里的损失函数可输入多个参数，但是我其实只需要两个：y_true，y_predict），我这里一共定义了两个损失函数，分别是Lap_Loss和SSIM_Loss. # 第一个损失函数 class Lap_Loss(layers.Layer): def __init__(self, **kwargs): super(Lap_Loss, self).__init__(**kwargs) def LapPyramid(self, input): # 构建拉普拉斯金字塔，进行下采样: 步长为2，先下采样，后上采样，再做差值 down0 = tf.image.resize(input, [128, 128], preserve_aspect_ratio=True, method="gaussian") up0 = tf.image.resize(down0, [256, 256],preserve_aspect_ratio=True, method="gaussian") diff0 = Subtract()([input,up0]) # 和原图的差 down1 = tf.image.resize(down0, [64, 64],preserve_aspect_ratio=True, method="gaussian") up1 = tf.image.resize(down1, [128, 128],preserve_aspect_ratio=True, method="gaussian") diff1 = Subtract()([down0, up1]) # 和原图/2的差 down2 = tf.image.resize(down1, [32, 32],preserve_aspect_ratio=True, method="gaussian") up2 = tf.image.resize(down2, [64, 64],preserve_aspect_ratio=True, method="gaussian") diff2 = Subtract()([down1, up2]) # 和原图/4的差 return diff0,diff1,diff2 def adaptive_loss(self,t,p,scale): mse_loss = tf.abs(t-p) panduan = mse_loss > scale large = mse_loss+scale*10 result = tf.where(panduan, large, mse_loss) out = tf.reduce_mean(tf.square(result)) return out def call(self, inputs, **kwargs): """ # inputs：Input tensor, or list/tuple of input tensors. 如上，父类layers的call方法明确要求inputs为一个tensor，或者包含多个tensor的列表/元组 所以这里不能直接接受多个入参，需要把多个入参封装成列表/元组的形式然后在函数中自行解包，否则会报错。 """ # 解包入参 y_true, y_pred = inputs # 复杂的损失函数 t_diff0, t_diff1, t_diff2 = self.LapPyramid(y_true) p_diff0, p_diff1, p_diff2 = self.LapPyramid(y_pred) diff2_loss = 1 - tf.image.ssim(t_diff2, p_diff2,max_val=2) diff1_loss = self.adaptive_loss(t_diff1,p_diff1,0.05) diff0_loss = self.adaptive_loss(t_diff1, p_diff1, 0.02) Lap_loss = diff2_loss + 0.1*diff1_loss + 0.1*diff0_loss # 重点：把自定义的loss添加进层使其生效，同时加入metric方便在KERAS的进度条上实时追踪 self.add_loss(Lap_loss , inputs=True) self.add_metric(Lap_loss , aggregation="mean", name="Lap_loss") return Lap_loss # 第二个损失函数 class SSIM_image_Loss(layers.Layer): def __init__(self, **kwargs): super(SSIM_image_Loss, self).__init__(**kwargs) def call(self, inputs, **kwargs): y_true, y_pred = inputs ssim_loss = 1 - tf.image.ssim(y_true, y_pred,max_val=2) img_ssim = ssim(y_true,y_pred) self.add_loss(ssim_loss, inputs=True) self.add_metric(img_ssim, aggregation="mean", name="ssim") return ssim_loss在定义完需要的损失函数后就可以构建模型了。 def comb_model(): # 这里需要给Input定义名字，后面构建数据输入需要用到 lr = Input(shape=(None, None, 3),name="lr_img") hr_image = Input(shape=(None, None, 3),name="hr_img") #one_gpu_model是之前写好的超分模型 sr_image = one_gpu_model(lr) loss_1 = PyramidLoss()([hr_image,sr_image]) loss_2 = SSIM_image_Loss()([hr_image,sr_image]) model = Model(inputs=[lr, hr_image], outputs=[sr_image, loss_1, loss_2]) model.compile(optimizer=Adam(lr=params['lr_init'])) return model由于我已经将自定义的两个loss函数以layer的方式加入到model中，因此在complie的时候我只对optimizer做了指定。可以看到，现在model的input=[lr, hr_image]，那么在进行训练时就需要用到：model.fit(x=input_data, y=None),因此需要对输入的数据进行更改。下面对更改前的数据生成代码和更改后的代码进行对比。（如果对tf的数据读入不太清楚的话，可以参考我之前的文章tensorflow数据读入的基础步骤） # 更改前的数据生成代码 def get_dataset(lr_path, hr_path, ext): lr_sorted_paths = get_sorted_image_path(lr_path, ext) hr_sorted_paths = get_sorted_image_path(hr_path, ext) # 打包hr和lr的所有地址 并进行shuffle打乱 lr_hr_sorted_paths = list(zip(lr_sorted_paths[:], hr_sorted_paths[:])) random.shuffle(lr_hr_sorted_paths) lr_sorted_paths, hr_sorted_paths = zip(*lr_hr_sorted_paths) # 将hr和lr组合成元组形式 ds = tf.data.Dataset.from_tensor_slices((list(lr_sorted_paths), list(hr_sorted_paths))) def load_and_preprocess_lr_hr_images(lr_path, hr_path, ext=ext): return load_and_preprocess_image(lr_path, ext), load_and_preprocess_image(hr_path, ext) lr_hr_ds = ds.map(load_and_preprocess_lr_hr_images, num_parallel_calls=4) return lr_hr_ds, len(lr_sorted_paths) # 后续的.shuffle()，.batch()等操作在此省略了此时，如果使用model.fit(lr_hr_ds)进行训练，keras的内部机制会直接让x=lr,y=hr，这样就无法满足input=[lr, hr_image]的要求。那么下面我们来看更改后的代码： def load_and_preprocess_lr_hr_images(lr_path, hr_path, ext=ext): lr = load_and_preprocess_image(lr_path, ext) hr = load_and_preprocess_image(hr_path, ext) inputs = {"lr_img":lr,"hr_img":hr} targets = {} return inputs, targets lr_hr_ds = ds.map(load_and_preprocess_lr_hr_images, num_parallel_calls=4)代码变动的地方只有load_and_preprocess_lr_hr_images这个函数。这里采用字典的形式，将lr和hr放入同一个字典中，需要注意字典中key的命名是要和model中Input的命名一致的。lr = Input(shape=(None, None, 3),name="lr_img") hr_image = Input(shape=(None, None, 3),name="hr_img")。这时运行model.fit(x=lr_hr_ds, y=None)就能够满足input的多输入要求了。 下面就是熟悉的keras进度条啦，因为我定义了两个损失函数，这里的loss其实是Lap_Loss和SSIM_Loss之和。 参考文章 Tensorflow2.0中复杂损失函数实现 |
虽然使用Layer.Subclass的写法看起来代码量会比较多，但是代码的完整度会高很多，后续用起来也更方便。【本文地址】